from Parameters import *
from PolymorphismTables import PolymorphismTables
from R2_Tools import R2
from SiteFrequencySpectrum import SiteFrequencySpectrum
import infiniteSitesTools as ist 
from Polyspect_Operator import Polyspect
from Databases import DB_Polyspect
from Databases import DB_Pluz

import math
import numpy.matlib as np
from scipy.stats import mstats
from scipy.stats import percentileofscore

import sys

def main(argv):
	
	print "***************************************************"
	print "********************* New run *********************"
	print "***************************************************"
	print ""
	
	tau = float(argv[0])
	
	mySFSs = None
	db = DB_Pluz2()

	#alphaToCheck = [5.0,10.0,15.0,20.0,25.0,30.0,35.0,40.0,45.0,50.0,55.0,60.0,65.0,70.0,75.0,80.0,85.0,90.0,95.0,100.0]
	alphaToCheck = np.arange(2.0, 100.0, 2.0)
	tauToCheck = [ tau ] #, 100, 200, 500, 1000, 2000, 5000, 10000]
	
	optCoord = np.ravel(np.zeros( 28 ) )
	
	#ccoord = ist.getOp_TajimasD(28)
	#ccoord = [0.000, -0.246, 0.045, -0.323, -0.022, -0.121, 0.053, -0.086, 0.032, 0.005, \
	#		  0.202, -0.196, 0.406, 0.148, 0.014, 0.069, 0.387, 0.314, 0.053, -0.234, \
	#		  0.010, -0.018, 0.179, 0.158, -0.321, 0.154, -0.063, -0.2050]
	
	ccoord = [0.000, -0.548, -0.336, -0.186, -0.080, 0.126, 0.023, -0.040, 0.162, \
			   0.092,  0.025,  0.221,  0.215,  0.058,-0.006,-0.156, -0.047,-0.180, \
			   0.083, -0.077, -0.100, -0.020,  0.204,-0.244, 0.117,  0.302, 0.119, 0.301]
	
	for i, c in enumerate( ccoord ):
		optCoord[i] = c

	biakaSFS = []
	manSFS = []
	sanSFS = []
	
	#1pMB4
	biakaSFS.append([0, 6,4,4,1,1,2,0,2,2,0,0,0,1,0,0,0,0,0,1,0,0,0,0,1,1,1,0])
	manSFS.append(  [0,11,4,4,7,4,0,2,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,2,0,1,0])
	sanSFS.append(  [0,14,0,2,4,1,2,1,2,0,0,0,0,1,2,0,0,0,0,0,0,0,0,0,0,0,0,0])

	#4qMB105
	biakaSFS.append([0, 5,5,3,0,1,1,0,0,0,0,0,1,3,0,0,0,0,0,0,1,0,0,0,0,0,0,0])
	manSFS.append(  [0, 1,2,5,5,0,0,4,1,1,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 3,5,6,5,2,4,1,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0])
	
	#4qMB181
	biakaSFS.append([0, 8,6,7,0,2,0,2,3,0,0,0,1,0,0,0,2,0,0,0,0,0,0,0,0,1,0,0])
	manSFS.append(  [0, 7,3,1,2,0,1,5,0,2,1,0,0,0,0,0,1,0,0,0,0,0,1,0,0,0,0,0])
	sanSFS.append(  [0, 5,3,2,1,3,2,0,1,3,0,0,0,0,0,2,0,0,0,0,0,0,0,0,0,0,0,0])

	#5pMB4
	biakaSFS.append([0, 6,3,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0])
	manSFS.append(  [0, 4,1,1,1,0,0,1,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 0,7,0,0,0,0,0,1,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0])

	#5pMB10
	biakaSFS.append([0,10,3,2,1,4,0,1,0,2,0,2,2,0,1,1,0,0,3,1,0,0,0,0,0,0,1,1])
	manSFS.append(  [0, 7,2,0,3,3,3,0,0,0,3,1,1,4,1,0,2,1,0,0,0,0,0,0,3,0,2,0])
	sanSFS.append(  [0, 4,3,2,1,1,0,0,0,1,0,1,0,4,0,0,0,0,0,0,0,0,0,0,0,0,0,0])

	#5qMB128
	biakaSFS.append([0,10,1,4,5,1,0,0,0,1,1,0,0,0,0,0,0,0,1,1,0,0,1,0,1,6,0,0])
	manSFS.append(  [0, 7,8,6,0,2,2,0,0,1,0,0,0,1,1,0,0,1,0,0,0,0,0,1,0,6,0,1])
	sanSFS.append(  [0, 6,4,0,0,1,1,1,0,0,0,0,0,0,0,1,0,1,0,0,0,0,0,0,0,0,0,0])

	#6pMB14
	biakaSFS.append([0, 3,9,3,7,2,1,6,3,0,0,0,0,0,0,3,0,0,0,0,1,6,0,0,1,0,0,0])
	manSFS.append(  [0, 8,4,5,3,4,2,1,2,1,7,1,1,0,1,0,0,1,4,1,0,0,0,0,0,2,0,0])
	sanSFS.append(  [0, 6,3,1,0,3,1,6,2,1,1,4,2,1,0,1,0,0,0,0,0,0,0,0,0,0,0,0])

	#6qMB164
	biakaSFS.append([0, 1,3,0,1,2,0,1,1,0,0,0,0,0,0,0,1,0,0,1,0,0,0,0,0,0,0,0])
	manSFS.append(  [0, 7,2,1,1,0,0,2,1,0,0,0,0,0,0,0,1,0,1,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 5,6,0,0,1,0,1,0,0,1,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0])

	#7pMB8
	biakaSFS.append([0,14,4,6,2,1,1,2,1,1,0,1,3,0,1,0,0,0,1,0,1,0,0,0,1,3,0,0])
	manSFS.append(  [0, 9,9,2,1,2,0,0,0,0,0,1,0,1,1,1,0,1,0,0,2,0,0,0,0,1,4,2])
	sanSFS.append(  [0,16,6,5,5,0,0,1,0,3,1,0,0,0,0,3,0,2,0,0,0,0,0,0,0,0,0,0])

	#8pMB5
	biakaSFS.append([0,14,5,8,2,0,1,1,1,0,0,0,0,3,0,0,1,1,0,1,1,1,0,0,0,1,0,0])
	manSFS.append(  [0, 7,3,1,0,6,1,2,0,1,0,1,1,3,0,0,0,0,0,0,0,2,0,0,0,0,0,1])
	sanSFS.append(  [0,10,6,7,4,2,2,0,2,0,0,0,1,1,0,2,1,0,0,0,0,0,0,0,0,0,0,0])

	#10qMB119
	biakaSFS.append([0, 6,2,1,4,1,1,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,0,0,0,0,0])
	manSFS.append(  [0, 5,1,1,1,0,1,0,0,2,0,1,1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 6,8,2,2,0,0,2,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0])

	#10qMB128
	biakaSFS.append([0, 7,4,4,0,0,0,1,1,3,1,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0])
	manSFS.append(  [0,10,1,1,0,0,6,3,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 6,5,4,0,0,0,1,1,2,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0])

	
	#12pMB46
	biakaSFS.append([0, 9,2,4,0,3,0,0,0,1,0,1,1,2,2,0,0,0,0,0,0,0,0,0,0,0,0,1])
	manSFS.append(  [0, 3,0,3,5,1,0,0,0,1,1,0,0,0,0,1,0,1,1,1,0,0,0,0,0,0,0,0]) 
	sanSFS.append(  [0, 8,5,1,1,2,0,3,2,1,0,0,0,1,1,0,0,1,0,0,0,0,0,0,0,0,0,0])

	#13qMB107
	biakaSFS.append([0, 6, 6,2,0,0,0,0,2,1,2,1,2,0,0,0,0,1,0,0,1,1,0,1,0,0,0,0])
	manSFS.append(  [0, 3, 0,0,2,1,0,0,0,0,1,0,0,0,2,1,1,0,1,1,0,0,0,2,1,0,0,0])
	sanSFS.append(  [0, 9,10,2,1,0,2,1,1,1,1,0,0,0,1,0,2,0,0,0,0,0,0,0,0,0,0,0])

	
	#13qMB108
	biakaSFS.append([0,12,3,3,2,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,1,0,0,0,2])
	manSFS.append(  [0, 5,1,1,4,4,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,1,0,0,0,1])
	sanSFS.append(  [0,10,0,4,0,0,0,1,1,1,0,0,0,0,0,2,0,0,0,0,0,0,0,0,0,0,0,0])

	#16pMB17
	biakaSFS.append([0, 8,6,3,5,2,2,0,0,2,0,0,0,0,0,1,0,1,0,1,2,0,0,0,2,1,0,0])
	manSFS.append(  [0,13,2,1,2,2,3,2,0,0,0,0,0,1,0,0,0,0,0,1,0,1,1,2,1,2,1,0])
	sanSFS.append(  [0, 5,3,4,1,1,1,0,0,0,0,1,0,0,0,1,1,2,0,0,0,0,0,0,1,1,0,1])

	#18pMB7
	biakaSFS.append([0, 4,3,0,1,1,2,0,0,2,0,1,0,1,0,0,0,0,0,0,0,1,0,0,0,0,0,0])
	manSFS.append(  [0,11,3,3,0,2,0,0,0,1,1,0,1,1,0,1,2,0,0,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 7,3,4,2,2,2,0,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0])

	#18qMB73
	biakaSFS.append([0,15,6,4,2,1,2,4,2,1,1,0,2,0,0,1,0,0,1,2,0,0,0,1,0,1,1,0])
	manSFS.append(  [0, 9,6,4,3,1,1,3,2,3,2,1,0,0,1,1,1,0,0,0,0,0,2,2,0,1,0,0])
	sanSFS.append(  [0,16,3,2,3,4,1,2,0,0,0,3,0,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0])

	#19qMB35
	biakaSFS.append([0,10,6,2,1,0,0,1,1,0,1,1,1,0,0,0,1,0,0,0,0,2,0,1,1,0,0,0])
	manSFS.append(  [0, 8,5,1,1,0,0,0,0,0,1,0,0,0,0,0,0,1,1,1,0,0,0,0,0,0,4,0])
	sanSFS.append(  [0,10,5,1,2,2,1,1,2,1,0,1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0])

	#20pMB7
	biakaSFS.append([0,11,1,1, 1,1,0,1,1,0,1,2,5,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0])
	manSFS.append(  [0, 7,0,1, 2,1,4,7,0,0,0,0,0,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0])
	sanSFS.append(  [0, 2,0,0,13,0,2,2,0,0,0,0,5,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0])

	biakaRegionStats = ist.operationOnAllSFS( optCoord, biakaSFS )
	biakaMean = mstats.tmean(biakaRegionStats)
	biakaVar = mstats.tvar(biakaRegionStats)
	
	print biakaRegionStats
	print biakaMean, biakaVar
	
	manRegionStats = ist.operationOnAllSFS( optCoord, manSFS )
	manMean = mstats.tmean(manRegionStats)
	manVar = mstats.tvar(manRegionStats)
	
	print manRegionStats
	print manMean, manVar

	sanRegionStats = ist.operationOnAllSFS( optCoord, sanSFS )
	sanMean = mstats.tmean(sanRegionStats)
	sanVar = mstats.tvar(sanRegionStats)
	
	print sanRegionStats
	print sanMean, sanVar
	
	
	for a, alpha in enumerate(alphaToCheck):
		for t, tau in enumerate(tauToCheck):

			summaryStatsforRegion = []
			
			for bSFS in biakaSFS:

				s = sum(bSFS)
				pop = getDefaultPopulationParameters(s)
				
				nullDemography = Demography( pop.N )
				nullDemography.addEpoch( pop.N, 0, True )
				print nullDemography
				print pop.getMSString_Null(), nullDemography.getMSString()
				
				
				N1 = pop.N * math.exp(-alpha * tau / (4.0*pop.N))
				
				altDemography = Demography(pop.N)
				altDemography.addEpoch( alpha, 0, False)
				altDemography.addEpoch( N1, tau, True )
					
				simParams = SimulationParameters( pop, nullDemography )
				print simParams
			
				
				pt = PolymorphismTables()
				msCommand = simParams.getMSString_Alternative( altDemography )
				msOut = pt.runMS(msCommand)
				pt.readMS(None, msOut)
				mySFSs = SiteFrequencySpectrum(pt)
	
				summaryStats = ist.operationOnAllSFS(optCoord, mySFSs)
				#print summaryStats
				
				summaryStatsforRegion.append( summaryStats )
			
			biakaRegionMeans = []
			biakaRegionVars = []
				
			for i in range(len(summaryStats)):
				rowStats = []
				for k in range(20):
					s = summaryStatsforRegion[k][i]
					rowStats.append(s)
				biakaRegionMeans.append( mstats.tmean(rowStats) )
				biakaRegionVars.append( mstats.tvar(rowStats) )

			meanP = ist.middlePvalue( biakaRegionMeans, biakaMean )
			varP = ist.middlePvalue( biakaRegionVars, biakaVar )
			
			db.addPoint( "biaka", tau, alpha, meanP, varP, pop.numPowerSimulations ) 

			print "biaka\t", tau, alpha, meanP, varP


			manRegionMeans = []
			manRegionVars = []
				
			for i in range(len(summaryStats)):
				rowStats = []
				for k in range(20):
					s = summaryStatsforRegion[k][i]
					rowStats.append(s)
				manRegionMeans.append( mstats.tmean(rowStats) )
				manRegionVars.append( mstats.tvar(rowStats) )

			meanP = ist.middlePvalue( manRegionMeans, manMean )
			varP = ist.middlePvalue( manRegionVars, manVar )
			
			db.addPoint( "mandenka", tau, alpha, meanP, varP, pop.numPowerSimulations ) 

			print "man\t", tau, alpha, meanP, varP



			sanRegionMeans = []
			sanRegionVars = []
				
			for i in range(len(summaryStats)):
				rowStats = []
				for k in range(20):
					s = summaryStatsforRegion[k][i]
					rowStats.append(s)
				sanRegionMeans.append( mstats.tmean(rowStats) )
				sanRegionVars.append( mstats.tvar(rowStats) )

			meanP = ist.middlePvalue( sanRegionMeans, sanMean )
			varP = ist.middlePvalue( sanRegionVars, sanVar )
			
			db.addPoint( "san", tau, alpha, meanP, varP, pop.numPowerSimulations ) 

			print "san\t", tau, alpha, meanP, varP
			

def getDefaultPopulationParameters(s=30):
	
	popParameters = Population()
	
	popParameters.numChromosomes = 28
	popParameters.numSegregatingSites = s

	popParameters.rateMutation = 1E-8
	popParameters.rateRecombination = 0.0
	popParameters.sequenceLength = 25000

	popParameters.N = 10000
	popParameters.numSimulations = 25000
	popParameters.numPowerSimulations = 10000

	return popParameters


if __name__ == "__main__":
	main(sys.argv[1:])
